import torch  
import torch.nn as nn  
import torch.optim as optim  
from torchvision import datasets, transforms  
from torch.utils.data import DataLoader, random_split  
from modelstruct import MLP
  
# 定义数据转换  
transform = transforms.Compose([  
    transforms.ToTensor(),  
    transforms.Normalize((0.1307,), (0.3081,))  
])  
  
# 加载MNIST数据集  
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)  
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)  
  
# 划分训练集为训练子集和验证子集  
train_size = int(0.8 * len(train_dataset))  
val_size = len(train_dataset) - train_size  
train_subset, val_subset = random_split(train_dataset, [train_size, val_size])  
  
# 创建数据加载器  
train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)  
val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)  
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)  
  
# 创建多层感知机模型
model = MLP()  
  
# 定义损失函数和优化器  
criterion = nn.CrossEntropyLoss()  
optimizer = optim.Adam(model.parameters(), lr=0.001)  
  
# 训练模型  
num_epochs = 20  
for epoch in range(num_epochs):  
    model.train()  # 设置模型为训练模式  
    running_loss = 0.0  
    correct = 0  
    total = 0  
      
    # 遍历训练数据加载器  
    for images, labels in train_loader:  
        # 清空梯度  
        optimizer.zero_grad()            
        # 前向传播  
        outputs = model(images)           
        # 计算损失  
        loss = criterion(outputs, labels)      
        # 反向传播  
        loss.backward()  
        # 更新权重  
        optimizer.step()      
        # 累计损失  
        running_loss += loss.item() * images.size(0)  
        # 计算预测正确的样本数  
        _, predicted = torch.max(outputs, 1)  
        total += labels.size(0)  
        correct += (predicted == labels).sum().item()  
      
    # 计算平均损失和准确率  
    epoch_loss = running_loss / total  
    epoch_acc = 100 * correct / total  
      
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}, Acc: {epoch_acc:.2f}%')  
  
import torch  
  
# 验证阶段  
model.eval()  # 设置为评估模式  
val_loss = 0.0  
val_acc = 0.0  
  
with torch.no_grad():  
    for images, labels in val_loader:  
        images, labels = images.to(device), labels.to(device) 
        outputs = model(images)  
        loss = criterion(outputs, labels)  

        val_loss += loss.item() * images.size(0)  
        _, predicted = torch.max(outputs, 1)  
        val_acc += (predicted == labels).sum().item()  
  
val_loss /= len(val_dataset)  
val_acc /= len(val_dataset)  
  
print(f"Epoch [{epoch+1}/{num_epochs}], Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")  
  
# 在测试集上评估模型  
model.eval()  
correct = 0  
total = 0  
  
with torch.no_grad():  
    for images, labels in test_loader:  
        images, labels = images.to(device), labels.to(device)  
        outputs = model(images)  
        _, predicted = torch.max(outputs, 1)  
        total += labels.size(0)  
        correct += (predicted == labels).sum().item()  
  
accuracy = correct / total  
print(f"Test Accuracy: {accuracy:.4f}")  
  
# 保存模型参数  
model_params = 'mnist_mlp.pth'  
torch.save(model.state_dict(), model_params)  
print(f"PyTorch MLP多层感知机模型参数已保存到: {model_params}")